package dict

import (
	"abdis/interface/datastruct"
	"abdis/lib/utils"
	"abdis/lib/wildcard"
	"math"
	"sync"
	"sync/atomic"
)

// shard 是一片存储空间
type shard struct {
	m     map[string]interface{}
	mutex sync.RWMutex
}

// ConcurrentDict 线程安全的存储器
type ConcurrentDict struct {
	table      []*shard
	count      int32
	shardCount int
}

func (dict *ConcurrentDict) deltaCount(delta int32) int32 {
	return atomic.AddInt32(&dict.count, delta)
}

func (dict *ConcurrentDict) checkValid() {
	if dict == nil {
		panic("dict is nil")
	}
}

// 计算开启的存储片数，为2的幂次方
func computeCapacity(param int) (size int) {
	if param <= 16 {
		return 16
	}

	n := param - 1
	n |= n >> 1
	n |= n >> 2
	n |= n >> 4
	n |= n >> 8
	n |= n >> 16
	if n < 0 {
		return math.MaxInt32
	}
	return n + 1
}

// MakeConcurrent 生成存储空间，实现内存数据库的存储空间
func MakeConcurrent(shardCount int) *ConcurrentDict {
	shardCount = computeCapacity(shardCount)
	table := make([]*shard, shardCount)
	for i := 0; i < shardCount; i++ {
		table[i] = &shard{
			m: make(map[string]interface{}),
		}
	}
	return &ConcurrentDict{
		table:      table,
		count:      0,
		shardCount: shardCount,
	}
}

// spread 找到该 hash 值应该存放的片数序号
func (dict *ConcurrentDict) spread(hash uint32) uint32 {
	dict.checkValid()
	tableSize := uint32(len(dict.table))
	// tableSize - 1 的二进制全1， & 操作使得结果在范围内
	return (tableSize - 1) & hash
}

func (dict *ConcurrentDict) getShard(index uint32) *shard {
	dict.checkValid()
	return dict.table[index]
}

func (dict *ConcurrentDict) Get(key string) (val interface{}, exists bool) {
	dict.checkValid()
	hash := utils.Fnv32(key)
	index := dict.spread(hash)
	s := dict.getShard(index)
	// 读取操作，加上读锁
	s.mutex.RLock()
	defer s.mutex.RUnlock()
	val, exists = s.m[key]
	return
}

func (dict *ConcurrentDict) GetWithLock(key string) (val interface{}, exists bool) {
	dict.checkValid()
	hash := utils.Fnv32(key)
	index := dict.spread(hash)
	s := dict.getShard(index)
	val, exists = s.m[key]
	return
}

// Put 将数据放入内存数据库中
func (dict *ConcurrentDict) Put(key string, val interface{}) int {
	dict.checkValid()
	hash := utils.Fnv32(key)
	index := dict.spread(hash)
	s := dict.getShard(index)
	// 写操作，上写锁
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if _, ok := s.m[key]; ok {
		// 原本有数据，新增 0
		s.m[key] = val
		return 0
	}
	dict.deltaCount(1)
	s.m[key] = val
	return 1
}

func (dict *ConcurrentDict) PutWithLock(key string, val interface{}) int {
	dict.checkValid()
	hash := utils.Fnv32(key)
	index := dict.spread(hash)
	s := dict.getShard(index)

	if _, ok := s.m[key]; ok {
		// 原本有数据，新增 0
		s.m[key] = val
		return 0
	}
	dict.deltaCount(1)
	s.m[key] = val
	return 1
}

func (dict *ConcurrentDict) PutIfAbsentWithLock(key string, val any) int {
	dict.checkValid()
	s := dict.getShard(dict.spread(utils.Fnv32(key)))

	if _, ok := s.m[key]; ok {
		return 0
	}

	s.m[key] = val
	dict.deltaCount(1)
	return 1
}

// Len 返回数据量的条数
func (dict *ConcurrentDict) Len() int {
	dict.checkValid()
	return int(atomic.LoadInt32(&dict.count))
}

// RWLocks locks write keys and read keys together. allow duplicate keys
func (dict *ConcurrentDict) RWLocks(writeKeys []string, readKeys []string) {
	keys := append(writeKeys, readKeys...)
	indices := utils.ToLockIndices(keys, false, dict.spread)
	writeIndexSet := make(map[uint32]struct{})
	for _, wKey := range writeKeys {
		idx := dict.spread(utils.Fnv32(wKey))
		writeIndexSet[idx] = struct{}{}
	}
	for _, index := range indices {
		_, w := writeIndexSet[index]
		m := &dict.table[index].mutex
		if w {
			m.Lock()
		} else {
			m.RLock()
		}
	}
}

func (dict *ConcurrentDict) RWUnLocks(writeKeys []string, readKeys []string) {
	keys := append(writeKeys, readKeys...)
	indices := utils.ToLockIndices(keys, true, dict.spread)
	writeIndexSet := make(map[uint32]struct{})
	for _, wKey := range writeKeys {
		idx := dict.spread(utils.Fnv32(wKey))
		writeIndexSet[idx] = struct{}{}
	}
	for _, index := range indices {
		_, w := writeIndexSet[index]
		m := &dict.table[index].mutex
		if w {
			m.Unlock()
		} else {
			m.RUnlock()
		}
	}
}

func (dict *ConcurrentDict) Remove(key string) (val any, result int) {
	dict.checkValid()
	s := dict.getShard(dict.spread(utils.Fnv32(key)))

	s.mutex.Lock()
	defer s.mutex.Unlock()

	if val, ok := s.m[key]; ok {
		delete(s.m, key)
		dict.deltaCount(-1)
		return val, 1
	}
	return val, 0
}

func (dict *ConcurrentDict) RemoveWithLock(key string) (val interface{}, result int) {
	dict.checkValid()
	s := dict.getShard(dict.spread(utils.Fnv32(key)))

	if val, ok := s.m[key]; ok {
		delete(s.m, key)
		dict.deltaCount(-1)
		return val, 1
	}
	return val, 0
}

func (dict *ConcurrentDict) ForEach(consumer datastruct.Consumer) {
	dict.checkValid()

	for _, s := range dict.table {
		s.mutex.RLock()
		fun := func() bool {
			defer s.mutex.RUnlock()
			for key, value := range s.m {
				if !consumer(key, value) {
					return false
				}
			}
			return true
		}
		if !fun() {
			break
		}
	}
}

func (dict *ConcurrentDict) Keys() []string {
	keys := make([]string, dict.Len())
	i := 0
	dict.ForEach(func(key string, val any) bool {
		keys[i] = key
		i++
		return true
	})
	return keys
}

func (dict *ConcurrentDict) Scan(cursor int, count int, pattern string) ([][]byte, int) {
	size := dict.Len()
	result := make([][]byte, 0)
	if pattern == "*" && count >= size {
		return utils.StringsToBytes(dict.Keys()), 0
	}

	match, err := wildcard.CompilePattern(pattern)
	if err != nil {
		return result, -1
	}

	shardCount := len(dict.table)
	shardIndex := cursor // 作为游标

	for shardIndex < shardCount {
		s := dict.table[shardIndex]
		s.mutex.RLock()

		if len(result)+len(s.m) > count && shardIndex > cursor {
			s.mutex.RUnlock()
			return result, shardIndex
		}

		for key := range s.m {
			if pattern == "*" || match.IsMatch(key) {
				result = append(result, []byte(key))
			}
		}
		s.mutex.RUnlock()
		shardIndex++
	}

	return result, 0
}
